#include "stdafx.h"
#include "shmenu.h"

namespace
{
	struct tagTrackHelpObj
	{
		HWND m_hWnd;
		WNDPROC oldProc;
		unsigned char * execode;
		IContextMenu3 * icm3;

		template<class T>
		DWORD_PTR todp(T obj)
		{
			typedef int foo[sizeof(obj) == sizeof(DWORD_PTR) ? 1 : -1];
			return *(DWORD_PTR*)&obj;
		}

		tagTrackHelpObj(HWND hWnd,IContextMenu3 * icm3)
		{
			this->icm3 = icm3;
			m_hWnd = hWnd;
			execode = 0;
			oldProc = 0;

			if(icm3 && IsWindow(hWnd))
			{
				execode = new unsigned char [64];
				VirtualProtect(execode,64,PAGE_EXECUTE_READWRITE,0);
				/*
				00401000 >    B9 78563412            mov ecx,12345678
				00401009    - E9 F20F1000            jmp 00502000
				*/
				//*execode ++ = 0xcc;
				execode[0] = 0xb9;
				execode[5] = 0xe9;
				*(void**)&execode[1] = this;
				*(size_t*)&execode[6] = todp(&tagTrackHelpObj::WinProc) - (DWORD_PTR)(execode + 10);
				oldProc = (WNDPROC)SetWindowLongPtr(m_hWnd,GWL_WNDPROC,(LONG_PTR)execode);
			}
		}
		void Detach()
		{
			if(oldProc) SetWindowLongPtr(m_hWnd,GWL_WNDPROC,(LONG_PTR)oldProc);
			delete [] execode;
			execode = 0;
			oldProc = 0;
			m_hWnd = 0;
		}
		~tagTrackHelpObj()
		{
			Detach();
		}

		LRESULT WinProc(HWND hWnd,UINT message,WPARAM wParam,LPARAM lParam)
		{
			LRESULT res;

			switch (message)
			{
			case WM_MENUCHAR:	// only supported by IContextMenu3
				icm3->HandleMenuMsg2 (message, wParam, lParam, &res);
				return res;
			case WM_DRAWITEM:
			case WM_MEASUREITEM:
				if (wParam) 
					break; // if wParam != 0 then the message is not menu-related
			case WM_INITMENUPOPUP:
				icm3->HandleMenuMsg (message, wParam, lParam);
				return (message == WM_INITMENUPOPUP ? 0 : TRUE); // inform caller that we handled WM_INITPOPUPMENU by ourself
			default:
				break;
			}
			// call original WndProc of window to prevent undefined bevhaviour of window
			return ::CallWindowProc (oldProc, hWnd, message, wParam, lParam);
		}
	};
}

CSimpleContextMenu::CSimpleContextMenu(HWND hWnd, LPCTSTR path)
{
	m_fPath = path;
	m_fPath.Replace('/', '\\');
	CComPtr<IShellFolder> fd, par;
	CComPtr<IContextMenu> icm;
	CComPtr<IContextMenu3> icm3;
	SHGetDesktopFolder(&fd);

	LPITEMIDLIST idl=0;
	LPCITEMIDLIST idlrela=0;
	fd->ParseDisplayName(hWnd, NULL, (LPOLESTR)(LPCTSTR)m_fPath, NULL, &idl, NULL);
	SHBindToParent(idl, IID_IShellFolder, (void**)&par, &idlrela);
	if (par && idlrela && idl)
	{
		par->GetUIObjectOf(hWnd, 1, &idlrela, IID_IContextMenu, 0, (void**)&icm);
		if (icm) icm->QueryInterface(IID_IContextMenu3, (void**)&m_icm3);
	}
	ILFree(idl);
}

HRESULT CSimpleContextMenu::Track(int x, int y, HWND hWnd)
{
	HMENU menu = ::CreatePopupMenu();
	if (!menu) return E_FAIL;

	DWORD dwFlags = GetFileAttributes(m_fPath);
	PreSetMenu(menu);
	if (m_icm3) m_icm3->QueryContextMenu(menu, GetMenuItemCount(menu), min_id,max_id,CMF_NORMAL);
	PostSetMenu(menu);

	DWORD pid=0, tid=0;
	int id = 0;
	tid = GetWindowThreadProcessId(hWnd, &pid);
	if (m_icm3 && pid == GetCurrentProcessId() && tid == GetCurrentThreadId())
	{
		tagTrackHelpObj to(hWnd, m_icm3);
		id = TrackPopupMenu(menu,TPM_RETURNCMD|TPM_LEFTALIGN, x, y, 0, hWnd, NULL);
	}
	else
	{
		id = TrackPopupMenu(menu,TPM_RETURNCMD|TPM_LEFTALIGN, x, y, 0, hWnd, NULL);
	}
	DestroyMenu(menu);

	if (id >= min_id && id <= max_id)
	{
		CMINVOKECOMMANDINFO cmi = {sizeof (CMINVOKECOMMANDINFO)};
		cmi.lpVerb = (LPSTR)MAKEINTRESOURCE(id - min_id);
		cmi.nShow = SW_SHOWNORMAL;
		return m_icm3->InvokeCommand(&cmi);
	}
	else
	{
		return ExecuteCmd(id, hWnd);
	}
}


//////////////////////////////////////////////////////////////////////////
//////////////////////////////////////////////////////////////////////////
//////////////////////////////////////////////////////////////////////////

void CContextMenuExt::CopyText(LPCTSTR text, HWND hWnd) const
{
	size_t sz = (_tcslen(text)+1)*sizeof(TCHAR);
	HANDLE gd = GlobalAlloc(GHND, sz);
	memcpy(GlobalLock(gd), text, sz);
	GlobalUnlock(gd);

	if (OpenClipboard(hWnd))
	{
		if (EmptyClipboard() && SetClipboardData(CF_UNICODETEXT, gd))
		{
			gd = 0;
		}
		CloseClipboard();
	}
	if (gd)GlobalFree(gd);
}

void CContextMenuExt::PostSetMenu(HMENU menu)
{
	int pos = 0;

	InsertMenu(menu, pos++,MF_BYPOSITION,max_id+1, _T("Copy &Name"));
	InsertMenu(menu, pos++,MF_BYPOSITION,max_id+2, _T("C&MD to Folder or Containing Folder"));
	if (m_bIsFile)
	{
		InsertMenu(menu, pos++,MF_BYPOSITION,max_id+3, _T("Open Conta&ining Folder") );
		InsertMenu(menu, pos++,MF_BYPOSITION,max_id+4, _T("Copy Containing Folder Name") );
	}
	InsertMenu(menu, pos++,MF_BYPOSITION|MF_SEPARATOR, 0, 0);
	SetMenuDefaultItem(menu, max_id+1, TRUE);
}

HRESULT CContextMenuExt::ExecuteCmd(int id, HWND hWnd)
{
	CString fdPath;
	int ifnd = m_fPath.ReverseFind('\\');
	if (ifnd>0 && m_bIsFile) fdPath = m_fPath.Left(ifnd+1);
	else fdPath = m_fPath;

	switch (id-max_id)
	{
	case 1: CopyText(m_fPath, hWnd); break;
	case 4: CopyText(fdPath, hWnd); break;
	case 2:
		{
			CString cmd;
			cmd.GetEnvironmentVariable(L"ComSpec");
			ShellExecuteW(hWnd, L"open", cmd, 0, fdPath, SW_SHOWNORMAL);
		}
		break;
	case 3:
		if (GetFileAttributesW(m_fPath) != ~0)
		{
			CString cmdline;
			cmdline.Format(L"/n,/select,\"%s\"", m_fPath);
			ShellExecuteW(hWnd, 0, L"explorer", cmdline, 0, SW_SHOWNORMAL);
		}
		else
		{
			ShellExecuteW(hWnd, L"open", fdPath, 0, fdPath, SW_SHOWNORMAL);
		}
		break;
	default:
		return E_UNEXPECTED;
	}
	return S_OK;
}

CContextMenuExt::CContextMenuExt(HWND hWnd, LPCTSTR path) : CSimpleContextMenu(hWnd,path)
{
	DWORD dwFlags = GetFileAttributes(m_fPath);
	bool isFile = true;
	if (dwFlags == INVALID_FILE_ATTRIBUTES)
	{
		isFile = m_fPath.Right(1) != _T("\\");
	}
	else
	{
		isFile = (dwFlags & FILE_ATTRIBUTE_DIRECTORY) == 0;
	}
	m_bIsFile = isFile;
}
